import numpy as np
from common_function import dcm_from_euler, rotation_matrix


class UAVDynamics:
    def __init__(self, params):
        self.params = params
        self.w_current = params.ModelInit_Rads

    def get_angular_velocity_derivative(self, M, J, omega):
        return np.linalg.inv(J) @ (M - np.cross(omega, J @ omega))

    def get_linear_velocity_derivative(self, F, m):
        return np.clip(F / m, -0.4, 0.4)

    def Obtain_force_torque(self, w, R, Cm, Ct, Vb, Cd, wb, Cdm, Jrp):
        M_rctcm = np.array([
            [-np.sin(np.pi / 4) * R * Ct, np.sin(np.pi / 4) * R * Ct, np.sin(np.pi / 4) * R * Ct,
             -np.sin(np.pi / 4) * R * Ct],  # Roll torque
            [np.sin(np.pi / 4) * R * Ct, -np.sin(np.pi / 4) * R * Ct, np.sin(np.pi / 4) * R * Ct,
             -np.sin(np.pi / 4) * R * Ct],  # Pitch torque
            [Cm, Cm, -Cm, -Cm]  # Yaw torque
        ])

        Mp = np.dot(M_rctcm, w ** 2)  # Torque

        Fp = np.array([0, 0, -np.sum(Ct * (w ** 2))])  # Thrust
        Fd = -Cd * Vb * np.abs(Vb) * 0.5  # Aerodynamic force
        Md = -Cdm * wb * np.abs(wb)  # Aerodynamic moment

        # Gyro moment
        Ga = np.zeros(3)
        Ga[0] = Jrp * wb[1] * (w[0] + w[1] - w[2] - w[3])
        Ga[1] = Jrp * wb[0] * (-w[0] - w[1] + w[2] + w[3])

        return Fp, Fd, Mp, Md, Ga

    def motor_dynamics(self, w_current, w_target, T, dt):
        dw = (w_target - w_current) / T
        return w_current + dw * dt

    def update_dynamics(self, state, w_target, dt):
        velB, angEuler, rateB = state

        m = self.params.ModelParam_uavMass
        J = self.params.ModelParam_uavJ
        g = self.params.ModelParam_envGravityAcc
        uavType = self.params.ModelParam_uavType
        R = self.params.ModelParam_uavR
        Cm = self.params.ModelParam_rotorCm
        Ct = self.params.ModelParam_rotorCt
        Cd = self.params.ModelParam_uavCd
        Cdm = self.params.ModelParam_uavCCm
        Jrp = self.params.ModelParam_motorJm
        T = self.params.ModelParam_motorT

        self.w_current = self.motor_dynamics(self.w_current, w_target, T, dt)
        Fp, Fd, Mp, Md, Ga = self.Obtain_force_torque(self.w_current, R, Cm, Ct, velB, Cd, rateB, Cdm, Jrp)
        dcm = dcm_from_euler(angEuler)
        total_force = Fp + Fd + np.array([0, 0, m * g]).dot(dcm)
        total_torque = Mp + Md + Ga
        velB_dot = self.get_linear_velocity_derivative(total_force, m)
        rateB_dot = self.get_angular_velocity_derivative(total_torque, J, rateB)

        velB = np.clip(velB + velB_dot * dt, -3, 3)
        rateB = rateB+rateB_dot * dt
        # angEuler = angEuler + rateB * dt

        return velB, angEuler, rateB

